package com.example.idholder.jdbc;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;

/**
 * <pre>
 * CREATE TABLE `id_hold` (
 *   `id` int NOT NULL AUTO_INCREMENT,
 *   `hold_time` datetime DEFAULT '2000-01-01 00:00:00',
 *   `refresh_time` datetime DEFAULT '2000-01-01 00:00:00' ON UPDATE CURRENT_TIMESTAMP,
 *   `refresh_count` int NOT NULL DEFAULT '0',
 *   `remark` varchar(512) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NOT NULL DEFAULT '',
 *   PRIMARY KEY (`id`)
 * ) ENGINE=InnoDB AUTO_INCREMENT=145 DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci;
 * </pre>
 * @author fanhang
 */
public class MysqlJdbcOperation implements JdbcOperation {
    private static final Logger log = LoggerFactory.getLogger(MysqlJdbcOperation.class);
    private final DataSource dataSource;
    // SQL
    private final String selectAvailableIdListSql;
    private final String lockForUpdateSql;
    private final String updateLockedSql;
    private final String refreshSql;

    public MysqlJdbcOperation(DataSource dataSource, String tableName) {
        this.dataSource = dataSource;
        String availableWhere = "hold_key = '' and hold_time is null and refresh_time is null and refresh_count = 0";
        // lock
        String sql = " select id from {tableName} where {availableWhere} limit 10 ";
        this.selectAvailableIdListSql = sql.replace("{tableName}", tableName).replace("{availableWhere}", availableWhere);
        sql = " select id from {tableName} where id = ? and {availableWhere} for update skip locked ";
        this.lockForUpdateSql = sql.replace("{tableName}", tableName).replace("{availableWhere}", availableWhere);
        // update locked
        sql = " update {tableName} set hold_key = ?, hold_time = CURRENT_TIMESTAMP, refresh_time = CURRENT_TIMESTAMP where id = ? ";
        this.updateLockedSql = sql.replace("{tableName}", tableName);
        // heartbeat
        sql = " update {tableName} set refresh_time = CURRENT_TIMESTAMP, refresh_count = refresh_count + 1 where id = ? and hold_key = ? ";
        this.refreshSql = sql.replace("{tableName}", tableName);
    }

    @Override
    public int lock(String holdKey) {
        while (true) {
            List<Integer> idList;
            try {
                idList = this.selectAvailableIdList();
            } catch (SQLException e) {
                throw new RuntimeException("获取可用ID列表失败");
            }
            if (idList == null || idList.isEmpty()) {
                throw new IllegalStateException("无可用ID");
            }
            for (Integer id : idList) {
                try {
                    if (lockAndUpdate(id, holdKey)) {
                        return id;
                    }
                } catch (SQLException e) {
                    log.error("SQLException", e);
                    throw new RuntimeException(e);
                }
            }
        }
    }

    private boolean lockAndUpdate(Integer id, String holdKey) throws SQLException {
        try (Connection conn = dataSource.getConnection()) {
            conn.setAutoCommit(false);
            Integer holdId = lockIdHold(conn, id);
            if (holdId == null) {
                return false;
            }
            int result = updateLocked(conn, holdKey, id);
            if (result == 1) {
                conn.commit();
                return true;
            } else {
                conn.rollback();
            }
        } catch (Exception ignored) {
        }
        return false;
    }

    @Override
    public int refresh(int id, String holdKey) {
        try (Connection conn = dataSource.getConnection()) {
            try (PreparedStatement pstmt = conn.prepareStatement(refreshSql)) {
                pstmt.setInt(1, id);
                pstmt.setString(2, holdKey);
                return pstmt.executeUpdate();
            }
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }

    private List<Integer> selectAvailableIdList() throws SQLException {
        try (Connection conn = dataSource.getConnection()) {
            try (PreparedStatement pstmt = conn.prepareStatement(selectAvailableIdListSql)) {
                List<Integer> result = new ArrayList<>();
                try (ResultSet rs = pstmt.executeQuery()) {
                    while (rs.next()) {
                        result.add(rs.getInt(1));
                    }
                    return result;
                }
            }
        }
    }

    private Integer lockIdHold(Connection conn, Integer id) throws SQLException {
        try (PreparedStatement pstmt = conn.prepareStatement(lockForUpdateSql)) {
            pstmt.setInt(1, id);
            try (ResultSet rs = pstmt.executeQuery()) {
                if (rs.next()) {
                    return rs.getInt(1);
                }
            }
            return null;
        }
    }

    private int updateLocked(Connection conn, String holdKey, Integer id) throws SQLException {
        try (PreparedStatement pstmt = conn.prepareStatement(updateLockedSql)) {
            pstmt.setString(1, holdKey);
            pstmt.setInt(2, id);
            return pstmt.executeUpdate();
        }
    }
}
